import datetime
import json
from typing import List, Iterator, Dict, Any

import heapdict
from kafka import KafkaConsumer
from pyshark.packet.packet import Packet

from utils import SUPPORTED_VPN_PROTOCOLS


class UDPS(object):
    def __init__(self):
        self.field_names = []  # type: List[str]


class Flow:
    def __init__(self,
                 start_time: float = 0,
                 end_time: float = 0,
                 recv_time: datetime.datetime = None,
                 packets: List[Packet] = None,
                 src_ip: str = "", dst_ip: str = "",
                 tunnel_src_ip: str = "0.0.0.0", tunnel_dst_ip: str = "0.0.0.0",
                 transport_type: str = "",
                 application_protocols: List[str] = None,
                 stream_index: int = -1,
                 client_port: int = 0, server_port: int = 0,
                 tunnel_src_port: int = 0, tunnel_dst_port: int = 0,
                 c2s_packet_length: int = -1, c2s_packet_num: int = -1,
                 s2c_packet_length: int = -1, s2c_packet_num: int = -1,
                 bidirectional_packet_length: int = -1, bidirectional_packet_num: int = -1,
                 url: str = "", sni: str = "",):
        self.start_time = start_time
        self.end_time = end_time
        self.packets = packets
        self.src_ip = src_ip
        self.dst_ip = dst_ip
        self.tunnel_src_ip = tunnel_src_ip
        self.tunnel_dst_ip = tunnel_dst_ip
        self.flow_transport_type = transport_type
        self.vpn_protocols = application_protocols
        self.stream_index = stream_index
        self.src_port = client_port
        self.dst_port = server_port
        self.tunnel_src_port = tunnel_src_port
        self.tunnel_dst_port = tunnel_dst_port
        self.c2s_packet_length = c2s_packet_length
        self.c2s_packet_num = c2s_packet_num
        self.s2c_packet_length = s2c_packet_length
        self.s2c_packet_num = s2c_packet_num
        self.bidirectional_packet_length = bidirectional_packet_length
        self.bidirectional_packet_num = bidirectional_packet_num
        self.recv_time = recv_time
        self.url = url
        self.sni = sni
        self.udps = UDPS()

    def __str__(self) -> str:
        json_dict = {
            "recv_time": self.recv_time.strftime("%Y-%m-%d %H:%M:%S"),
            "start_time": self.start_time,
            "end_time": self.end_time,
            "server_ip": self.dst_ip,
            "client_ip": self.src_ip,
            "tunnel_server_ip": self.tunnel_dst_ip,
            "tunnel_client_ip": self.tunnel_src_ip,
            "tunnel_server_port": self.tunnel_dst_port,
            "tunnel_client_port": self.tunnel_src_port,
            "server_port": self.dst_port,
            "client_port": self.src_port,
            "c2s_pkt_num": self.c2s_packet_num,
            "c2s_byte_num": self.c2s_packet_length,
            "s2c_pkt_num": self.s2c_packet_num,
            "s2c_pkt_length": self.s2c_packet_length,
            "URL": self.url,
            "SNI": self.sni,
        }
        for field in self.udps.field_names:
            vpn_layer = getattr(self.udps, field)  # type: Dict[str, Any]
            for key, value in vpn_layer.items():
                json_dict[key] = value
        return json.dumps(json_dict)


class FiveTupleInfo:
    def __init__(self,
                 client_ip, server_ip,
                 client_port, server_port,
                 tunnel_client_ip, tunnel_server_ip,
                 tunnel_client_port, tunnel_server_port,
                 c2s_packet_length, c2s_packet_num, s2c_packet_length, s2c_packet_num,
                 bidirectional_packet_length, bidirectional_packet_num,
                 protocols, url, sni):
        self.client_ip = client_ip
        self.client_port = client_port
        self.server_ip = server_ip
        self.server_port = server_port
        self.tunnel_client_ip = tunnel_client_ip
        self.tunnel_server_ip = tunnel_server_ip
        self.tunnel_client_port = tunnel_client_port
        self.tunnel_server_port = tunnel_server_port
        self.c2s_packet_length = c2s_packet_length
        self.c2s_packet_num = c2s_packet_num
        self.s2c_packet_length = s2c_packet_length
        self.s2c_packet_num = s2c_packet_num
        self.bidirectional_packet_length = bidirectional_packet_length
        self.bidirectional_packet_num = bidirectional_packet_num
        self.protocols = protocols
        self.url = url
        self.sni = sni


def get_five_tuple_info(packets: List[Packet]) -> FiveTupleInfo:
    """ generate the five tuple info. used in get_flow_info function

    :param packets: all the packets in the flow
    :return: five tuple information, including (client_ip, server_ip, c2s_packet_length, c2s_packet_num,
    s2c_packet_length, s2c_packet_num, bidirectional_packet_length, bidirectional_packet_num, client_port, server_port,
    protocols)
    """
    client_ip = ""
    server_ip = ""
    tunnel_client_ip = ""
    tunnel_server_ip = ""

    c2s_packet_length = 0
    c2s_packet_num = 0

    s2c_packet_length = 0
    s2c_packet_num = 0

    bidirectional_packet_length = 0
    bidirectional_packet_num = 0

    client_port = 0
    server_port = 0
    tunnel_client_port = 0
    tunnel_server_port = 0

    protocols = []  # type: List[str]
    url = []
    sni = ""

    for packet in packets:
        c2s = None
        in_vpn_tunnel = False
        # we reverse it because we want to traverse the layers from top to down
        # this is because there might be two tcp/ip layers in the total VPN packets
        # only the most top layer can represent the true tcp/ip layer
        for layer in packet.layers:
            layer_name = layer.layer_name.lower()

            if layer_name in SUPPORTED_VPN_PROTOCOLS:
                in_vpn_tunnel = True

            if layer_name in SUPPORTED_VPN_PROTOCOLS and layer_name not in protocols:
                protocols.append(layer_name)
            if layer_name == "ip":

                # if client ip is not defined
                if not in_vpn_tunnel and client_ip == "":
                    client_ip = str(layer.src)
                    server_ip = str(layer.dst)
                elif in_vpn_tunnel and tunnel_client_ip == "":
                    tunnel_client_ip = str(layer.src)
                    tunnel_server_ip = str(layer.dst)

                # if the direction of the packet is not defined
                if c2s is None:
                    c2s = layer.src == client_ip
            # if the client's transport layer is tcp or udp, the packet has the port
            if layer_name == "tcp" or layer_name == "udp":
                if not in_vpn_tunnel and client_port == 0:
                    client_port = int(layer.srcport)
                    server_port = int(layer.dstport)
                elif in_vpn_tunnel and tunnel_client_port == 0:
                    tunnel_client_port = int(layer.srcport)
                    tunnel_server_port = int(layer.dstport)
            if layer_name == "http":
                try:
                    url.append(str(layer.request_full_uri))
                except AttributeError:
                    pass

        # define the direction of one packet
        if c2s is not None:
            if c2s:
                c2s_packet_num += 1
                c2s_packet_length += int(packet.length)
            else:
                s2c_packet_num += 1
                s2c_packet_length += int(packet.length)

        # no matter what direction the packet is, the packet belongs to the bidirectional packet stream
        bidirectional_packet_num += 1
        bidirectional_packet_length += int(packet.length)

    return FiveTupleInfo(
        client_ip=client_ip, server_ip=server_ip,
        client_port=client_port, server_port=server_port,
        tunnel_client_ip=tunnel_client_ip, tunnel_server_ip=tunnel_server_ip,
        tunnel_client_port=tunnel_client_port, tunnel_server_port=tunnel_server_port,
        c2s_packet_num=c2s_packet_num, s2c_packet_num=s2c_packet_num,
        c2s_packet_length=c2s_packet_length, s2c_packet_length=s2c_packet_length,
        bidirectional_packet_num=bidirectional_packet_num, bidirectional_packet_length=bidirectional_packet_length,
        protocols=protocols,
        url=url,
        sni=sni,
    )


def get_flow_info(stream_packets: List[Packet], stream_index: int, flow_type: str):
    """ generate a flow containing all the info

    :param stream_packets: all the packets in the stream
    :param stream_index: the index of the stream
    :param flow_type: type name of the flow
    :return: a Flow object
    """
    start_time = float(stream_packets[0].sniff_timestamp)
    end_time = float(stream_packets[-1].sniff_timestamp)

    five_tuple_info = get_five_tuple_info(stream_packets)

    recv_time = datetime.datetime.now()

    return Flow(
        start_time=start_time,
        end_time=end_time,
        src_ip=five_tuple_info.client_ip,
        dst_ip=five_tuple_info.server_ip,
        tunnel_src_ip=five_tuple_info.tunnel_client_ip,
        tunnel_dst_ip=five_tuple_info.tunnel_server_ip,
        packets=stream_packets,
        transport_type=flow_type,
        application_protocols=five_tuple_info.protocols,
        stream_index=stream_index,
        client_port=five_tuple_info.client_port,
        server_port=five_tuple_info.server_port,
        tunnel_src_port=five_tuple_info.tunnel_client_port,
        tunnel_dst_port=five_tuple_info.tunnel_server_port,
        c2s_packet_length=five_tuple_info.c2s_packet_length,
        c2s_packet_num=five_tuple_info.c2s_packet_num,
        s2c_packet_length=five_tuple_info.s2c_packet_length,
        s2c_packet_num=five_tuple_info.s2c_packet_num,
        bidirectional_packet_length=five_tuple_info.bidirectional_packet_length,
        bidirectional_packet_num=five_tuple_info.bidirectional_packet_num,
        recv_time=recv_time,
        url=five_tuple_info.url,
        sni=five_tuple_info.sni,
    )


def generate_stream_id_by_ip(packet: Packet) -> int:
    """ generate a stream index of the packet(in udp/gre/esp)

    :param packet: coming packet
    :return: stream index
    """
    # there is no port for GRE, so we can only use src ip and dst ip to generate a GRE stream id
    return (hash(packet.ip.src) + hash(packet.ip.dst)) % 100000


def pop_outdated_streams(stream_lru_packet_time: heapdict.heapdict,
                         streams: Dict[int, List[Packet]],
                         packet: Packet,
                         flow_type: str,
                         timeout: float) -> Iterator[Flow]:
    """ pop all the outdated streams in the streams dict by using heapdict

    :param stream_lru_packet_time: the least_recently_used streams. Realized by using heapdict
    :param streams: the stream packets dict. key is stream index and value is stream packets
    :param packet: coming packet
    :param flow_type: the type name of the flow
    :param timeout: timeout of the flow_type
    :return: all the outdated flows
    """
    # use heapdict to pop/peek the lru stream index and its time
    if len(stream_lru_packet_time) == 0:
        return

    packet_time = float(packet.sniff_timestamp)
    (stream_index, stream_lru_time) = stream_lru_packet_time.peekitem()
    while packet_time - stream_lru_time > timeout:
        stream_lru_packet_time.popitem()
        stream_packets = streams[stream_index]
        del streams[stream_index]

        yield get_flow_info(stream_packets, stream_index, flow_type)

        if len(stream_lru_packet_time) == 0:
            return

        (stream_index, stream_lru_time) = stream_lru_packet_time.peekitem()


def kafka_consume(topics: str, bootstrap_servers: str):
    consumer = KafkaConsumer(
        topics,
        bootstrap_servers=bootstrap_servers,
    )

    for message in consumer:
        print(message)
